#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
小米便签网页版导出工具（手动配置版）

本工具不需要自动下载ChromeDriver，直接使用用户手动下载的驱动程序。
特点：
- 完全手动配置，不依赖网络下载
- 只支持JSON和TXT格式导出
- 简单易用，配置明确
"""

import os
import json
import argparse
import time
import logging
from datetime import datetime
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, WebDriverException

# 配置日志
def setup_logger():
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler('mi_note_export.log'),
            logging.StreamHandler()
        ]
    )
    return logging.getLogger(__name__)

logger = setup_logger()

class MiNoteWebExporter:
    def __init__(self, driver_path, chrome_path=None, output_dir="./exports", export_format="json", headless=False, timeout=30):
        """初始化导出工具"""
        self.driver_path = driver_path
        self.chrome_path = chrome_path
        self.output_dir = output_dir
        self.export_format = export_format.lower()
        self.headless = headless
        self.timeout = timeout
        self.driver = None
        
        # 验证输出目录存在
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
            logger.info(f"创建输出目录: {self.output_dir}")
    
    def init_browser(self):
        """初始化Chrome浏览器"""
        logger.info("正在初始化Chrome浏览器...")
        
        chrome_options = Options()
        
        # 设置Chrome浏览器路径（如果提供）
        if self.chrome_path:
            logger.info(f"使用指定的Chrome浏览器路径: {self.chrome_path}")
            chrome_options.binary_location = self.chrome_path
        
        # 无头模式设置
        if self.headless:
            chrome_options.add_argument("--headless")
            chrome_options.add_argument("--disable-gpu")
            chrome_options.add_argument("--window-size=1920,1080")
        
        # 添加其他优化参数
        chrome_options.add_argument("--no-sandbox")
        chrome_options.add_argument("--disable-dev-shm-usage")
        chrome_options.add_argument("--start-maximized")
        
        # 创建WebDriver服务
        service = Service(executable_path=self.driver_path)
        
        try:
            # 初始化WebDriver
            self.driver = webdriver.Chrome(service=service, options=chrome_options)
            logger.info("Chrome浏览器初始化成功")
            return True
        except WebDriverException as e:
            logger.error(f"Chrome浏览器初始化失败: {str(e)}")
            logger.error("请检查ChromeDriver和Chrome浏览器版本是否匹配")
            return False
    
    def login(self):
        """登录小米云服务"""
        logger.info("请在打开的浏览器中手动登录小米云服务")
        try:
            # 打开小米云服务便签页面
            self.driver.get("https://i.mi.com/note")
            
            # 等待用户登录 - 延长至5分钟
            print("\n=== 登录提示 ===")
            print("1. 请在打开的Chrome浏览器中登录您的小米账号")
            print("2. 登录成功后，请确保您能看到便签列表页面")
            print("3. 程序会自动检测登录状态并开始导出")
            print("4. 您有5分钟时间完成登录操作\n")
            
            logger.info("请在浏览器中完成登录，登录成功后等待自动继续...")
            
            # 等待便签页面加载完成（通过检查特定元素）
            wait = WebDriverWait(self.driver, 300)  # 延长至5分钟
            wait.until(EC.presence_of_element_located((By.CLASS_NAME, "note-list-item")))
            
            logger.info("登录成功，已进入便签页面")
            return True
        except TimeoutException:
            logger.error("登录超时，请检查是否成功登录到小米云服务便签页面")
            print("\n错误: 登录超时！请确认：")
            print("1. 您是否成功登录了小米云服务？")
            print("2. 您是否进入了便签页面（https://i.mi.com/note）？")
            print("3. 页面上是否显示了您的便签列表？")
            return False
        except Exception as e:
            logger.error(f"登录过程中发生错误: {str(e)}")
            return False
    
    def export_notes(self):
        """导出便签数据"""
        logger.info("开始导出便签数据...")
        print("\n开始导出便签数据，请稍候...")
        
        try:
            # 等待便签列表加载完成
            wait = WebDriverWait(self.driver, self.timeout)
            wait.until(EC.presence_of_element_located((By.CLASS_NAME, "note-list-item")))
            
            # 获取所有便签元素
            note_elements = self.driver.find_elements(By.CLASS_NAME, "note-list-item")
            logger.info(f"找到 {len(note_elements)} 条便签")
            print(f"找到 {len(note_elements)} 条便签，正在逐个导出...")
            
            notes_data = []
            
            # 遍历每个便签并提取内容
            for i, note_element in enumerate(note_elements, 1):
                try:
                    # 点击便签打开详情
                    note_element.click()
                    time.sleep(1)  # 等待内容加载
                    
                    # 提取标题和内容
                    title = "无标题"
                    content = ""
                    
                    # 尝试获取标题（不同版本的UI可能有所不同）
                    try:
                        title_element = self.driver.find_element(By.CLASS_NAME, "note-editor-title")
                        if title_element.text.strip():
                            title = title_element.text.strip()
                    except:
                        try:
                            # 尝试其他可能的标题元素类名
                            title_element = self.driver.find_element(By.CSS_SELECTOR, ".title-input")
                            if title_element.text.strip():
                                title = title_element.text.strip()
                        except:
                            pass
                    
                    # 尝试获取内容
                    try:
                        content_element = self.driver.find_element(By.CLASS_NAME, "note-editor-content")
                        content = content_element.text.strip()
                    except:
                        try:
                            # 尝试其他可能的内容元素类名
                            content_element = self.driver.find_element(By.CSS_SELECTOR, ".content-input")
                            content = content_element.text.strip()
                        except:
                            pass
                    
                    # 提取便签信息
                    note_data = {
                        "id": i,
                        "title": title,
                        "content": content,
                        "export_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    }
                    
                    notes_data.append(note_data)
                    logger.info(f"已提取便签 {i}/{len(note_elements)}: {title}")
                    print(f"已提取便签 {i}/{len(note_elements)}: {title}")
                    
                except Exception as e:
                    logger.error(f"提取便签 {i} 时发生错误: {str(e)}")
                    print(f"警告: 提取便签 {i} 时发生错误，跳过此条")
                    continue
            
            # 保存导出的数据
            self._save_exported_data(notes_data)
            return True
        except Exception as e:
            logger.error(f"导出便签过程中发生错误: {str(e)}")
            return False
    
    def _save_exported_data(self, notes_data):
        """保存导出的数据到文件"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        if self.export_format == "json":
            # 保存为JSON格式
            file_path = os.path.join(self.output_dir, f"mi_notes_{timestamp}.json")
            with open(file_path, "w", encoding="utf-8") as f:
                json.dump(notes_data, f, ensure_ascii=False, indent=2)
            logger.info(f"已将 {len(notes_data)} 条便签导出为JSON格式: {file_path}")
            print(f"\n已将 {len(notes_data)} 条便签导出为JSON格式:")
            print(f"{file_path}")
            
        elif self.export_format == "txt":
            # 保存为TXT格式（每个便签一个文件）
            txt_dir = os.path.join(self.output_dir, f"mi_notes_txt_{timestamp}")
            if not os.path.exists(txt_dir):
                os.makedirs(txt_dir)
            
            for note in notes_data:
                # 生成安全的文件名
                safe_title = "".join([c for c in note["title"] if c.isalnum() or c in " -_"]).rstrip()
                if not safe_title:
                    safe_title = f"无标题_{note['id']}"
                
                file_path = os.path.join(txt_dir, f"{safe_title}.txt")
                
                # 避免文件名冲突
                counter = 1
                base_path = file_path
                while os.path.exists(file_path):
                    name, ext = os.path.splitext(base_path)
                    file_path = f"{name}_{counter}{ext}"
                    counter += 1
                
                # 写入文件
                with open(file_path, "w", encoding="utf-8") as f:
                    if note["title"] != "无标题":
                        f.write(f"{note['title']}\n\n")
                    f.write(note["content"])
                    f.write(f"\n\n导出时间: {note['export_time']}")
            
            logger.info(f"已将 {len(notes_data)} 条便签导出为TXT格式: {txt_dir}")
            print(f"\n已将 {len(notes_data)} 条便签导出为TXT格式（每个便签一个文件）:")
            print(f"{txt_dir}")
            
    def close(self):
        """关闭浏览器"""
        if self.driver:
            self.driver.quit()
            logger.info("浏览器已关闭")

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='小米便签网页版导出工具（手动配置版）')
    parser.add_argument('--driver_path', required=True, help='ChromeDriver的路径')
    parser.add_argument('--chrome_path', help='Chrome浏览器的路径（可选）')
    parser.add_argument('--output_dir', default='./exports', help='导出文件的保存目录')
    parser.add_argument('--format', default='json', choices=['json', 'txt'], help='导出格式')
    parser.add_argument('--headless', action='store_true', help='是否以无头模式运行')
    parser.add_argument('--timeout', type=int, default=30, help='操作超时时间（秒）')
    return parser.parse_args()

def main():
    print("===== 小米便签网页版导出工具（手动配置版）=====")
    print("本工具需要您手动下载并配置ChromeDriver")
    print("===========================================")
    
    # 解析命令行参数
    args = parse_args()
    
    logger.info(f"导出格式: {args.format}")
    logger.info(f"输出目录: {args.output_dir}")
    logger.info(f"无头模式: {args.headless}")
    logger.info(f"超时时间: {args.timeout}秒")
    logger.info(f"ChromeDriver路径: {args.driver_path}")
    if args.chrome_path:
        logger.info(f"Chrome浏览器路径: {args.chrome_path}")
    
    # 验证ChromeDriver是否存在
    if not os.path.exists(args.driver_path):
        logger.error(f"ChromeDriver文件不存在: {args.driver_path}")
        print("错误: ChromeDriver文件不存在，请检查路径是否正确")
        return 1
    
    # 验证Chrome浏览器是否存在（如果提供了路径）
    if args.chrome_path and not os.path.exists(args.chrome_path):
        logger.error(f"Chrome浏览器文件不存在: {args.chrome_path}")
        print("错误: Chrome浏览器文件不存在，请检查路径是否正确")
        return 1
    
    try:
        # 创建导出工具实例
        exporter = MiNoteWebExporter(
            driver_path=args.driver_path,
            chrome_path=args.chrome_path,
            output_dir=args.output_dir,
            export_format=args.format,
            headless=args.headless,
            timeout=args.timeout
        )
        
        # 初始化浏览器
        if not exporter.init_browser():
            print("错误: 浏览器初始化失败，请检查ChromeDriver和Chrome版本是否匹配")
            return 1
        
        # 登录
        if not exporter.login():
            print("错误: 登录失败，请检查是否成功登录到小米云服务")
            exporter.close()
            return 1
        
        # 导出便签
        if not exporter.export_notes():
            print("错误: 导出便签失败")
            exporter.close()
            return 1
        
        # 完成导出
        exporter.close()
        logger.info("导出完成！")
        print("\n导出成功！")
        print(f"导出文件保存在: {args.output_dir}")
        print("感谢使用小米便签网页版导出工具！")
        return 0
        
    except Exception as e:
        logger.error(f"程序运行失败: {str(e)}")
        print(f"错误: {str(e)}")
        return 1

if __name__ == "__main__":
    exit_code = main()
    os._exit(exit_code)